from typing import Any

import torch
from torch import Tensor


class RnnModule(torch.nn.Module):
    r"""

    结构表达式

    .. math::

        h_t = \tanh(W_{ih} x_t + b_{ih} + W_{hh} h_{(t-1)} + b_{hh})

        y_t = U * h_t

    """

    def _forward_unimplemented(self, *input: Any) -> None:
        pass

    def __init__(self, input_size: int):
        super(RnnModule, self).__init__()
        self.rnn = torch.nn.RNN(input_size=input_size, hidden_size=6, num_layers=1, nonlinearity='tanh')
        self.result = torch.nn.Linear(6, 3)

    def forward(self, x: Tensor, batch_sizes: int):
        output, _ = self.rnn(x, torch.zeros((1, batch_sizes, 6)))
        output = torch.nn.utils.rnn.pad_packed_sequence(sequence=output)
        # batch维度放在第一位，使得后面的全连接计算符合习惯
        output = torch.transpose(output[0], 0, 1)
        y = self.result(output)
        return y
